import os
import cv2
import face_recognition
import tkinter as tk
from tkinter import filedialog, messagebox
import pickle

# 支持的图片文件扩展名
IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp')
# 存储已知人脸的编码信息
ENCODINGS_FILE = "known_faces_encodings.pkl"

# 加载已知图片并编码
def load_known_faces(known_faces_dir):
    known_faces = []
    known_names = []
    for filename in os.listdir(known_faces_dir):
        if not filename.lower().endswith(IMAGE_EXTENSIONS):
            continue
        filepath = os.path.join(known_faces_dir, filename)
        image = face_recognition.load_image_file(filepath)
        encoding = face_recognition.face_encodings(image)
        if encoding:
            known_faces.append(encoding[0])
            known_names.append(os.path.splitext(filename)[0])
    return known_faces, known_names

# 保存已知人脸编码
def save_encodings(known_faces, known_names, encodings_file):
    with open(encodings_file, 'wb') as f:
        pickle.dump((known_faces, known_names), f)
# 加载已知人脸编码
def load_encodings(encodings_file):
    with open(encodings_file, 'rb') as f:
        known_faces, known_names = pickle.load(f)
    return known_faces, known_names

# 检查图库是否有变动
def check_known_faces_dir(known_faces_dir):
    return sorted(os.listdir(known_faces_dir))

# 保存新面孔图片
def save_new_face(image_path, name, save_dir):
    print(f"保存新面孔: {name}.jpg")
    image = face_recognition.load_image_file(image_path)
    if image is None:
        print(f"无法读取图像: {image_path}")
        return
    filepath = os.path.join(save_dir, f"{name}.jpg")
    cv2.imwrite(filepath, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
# 绘制标签
def draw_labels(frame, face_locations, face_names):
    for (top, right, bottom, left), name in zip(face_locations, face_names):
        # 绘制矩形框
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
        
        # 获取人脸特征点
        landmarks = face_recognition.face_landmarks(frame[top:bottom, left:right])
        
        # 绘制特征点
        for face_landmarks in landmarks:
            for facial_feature in face_landmarks.keys():
                for point in face_landmarks[facial_feature]:
                    x, y = point
                    cv2.circle(frame, (left + x, top + y), 2, (0, 255, 0), -1)

        
        # 绘制标签
        cv2.rectangle(frame, (left, bottom - 25), (right, bottom), (255, 255, 255), cv2.FILLED)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (0, 0, 0), 2)

# 人脸识别函数
def recognize_faces(image_path, known_faces, known_names):
    image = face_recognition.load_image_file(image_path)
    face_locations = face_recognition.face_locations(image)
    
    if not face_locations:
        messagebox.showwarning("警告", "未检测到人脸，请重新选择图片。")
        return
    face_encodings = face_recognition.face_encodings(image, face_locations)

    face_names = []
    for face_encoding in face_encodings:
        matches = face_recognition.compare_faces(known_faces, face_encoding, tolerance=0.4)
        name = "Unknown"
        if True in matches:
            match_index = matches.index(True)
            name = known_names[match_index]
        face_names.append(name)

    frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    draw_labels(frame, face_locations, face_names)
    
    WINDOW_WIDTH = 800
    WINDOW_HEIGHT = 800
    frame_height, frame_width, _ = frame.shape
    aspect_ratio = frame_width / frame_height
    new_width = WINDOW_WIDTH
    new_height = int(new_width / aspect_ratio)
    if new_height > WINDOW_HEIGHT:
        new_height = WINDOW_HEIGHT
        new_width = int(new_height * aspect_ratio)
    frame = cv2.resize(frame, (new_width, new_height))
    
    cv2.imshow("Face Recognition", frame)
    cv2.waitKey(0)

# 选择图片的函数
def select_image():
    file_path = filedialog.askopenfilename()
    if file_path:
        recognize_faces(file_path, known_faces, known_names)

# 自定义名字输入框
def custom_input_dialog(prompt, title="输入名字", width=400, height=200):
    root = tk.Toplevel()
    root.title(title)
    root.geometry(f"{width}x{height}")
    
    tk.Label(root, text=prompt).pack(pady=10)
    
    name_var = tk.StringVar()
    entry = tk.Entry(root, textvariable=name_var)
    entry.pack(pady=10)
    entry.focus_set()
    
    def on_ok():
        root.destroy()
        root.name = name_var.get()
    
    def on_cancel():
        root.destroy()
        root.name = None
    
    ok_button = tk.Button(root, text="确定", command=on_ok)
    ok_button.pack(side="left", padx=100)
    
    cancel_button = tk.Button(root, text="取消", command=on_cancel)
    cancel_button.pack(side="left", padx=30)
    
    root.name = None
    root.transient()
    root.grab_set()
    root.wait_window()
    
    return root.name

# 添加新面孔的函数
def add_new_face():
    def on_select_image():
        file_path = filedialog.askopenfilename()
        if file_path:
            name = custom_input_dialog("请输入图片对应的名字:")
            if not name:
                messagebox.showerror("错误", "请输入名字")
                return
            save_new_face(file_path, name, known_faces_dir)
            
            image = face_recognition.load_image_file(file_path)
            encoding = face_recognition.face_encodings(image)
            if encoding:
                known_faces.append(encoding[0])
                known_names.append(name)
            
            save_encodings(known_faces, known_names, ENCODINGS_FILE)
            messagebox.showinfo("成功", f"{name} 已成功添加到已知人脸库")
    # 创建选择图片的窗口 
    add_image_window = tk.Toplevel()
    add_image_window.title("添加新面孔")
    window_width = 500
    window_height = 400
    add_image_window.geometry(f"{window_width}x{window_height}")
    select_button = tk.Button(add_image_window, text="选择图片", command=on_select_image)
    select_button.pack(pady=10)

# GUI界面
def create_gui():
    root = tk.Tk()
    root.title("人脸识别应用")
    root.geometry("500x400+400+300")

    select_button = tk.Button(root, text="选择需要识别的图片", command=select_image)
    select_button.pack(pady=20)

    add_button = tk.Button(root, text="添加新面孔", command=add_new_face)
    add_button.pack(pady=20)

    root.mainloop()

# 初始化已知人脸数据
known_faces_dir = "../step4/known_image"
current_faces_list = check_known_faces_dir(known_faces_dir)

# 检查是否已存在人脸编码文件(ENCODINGS_FILE)和人脸图片列表文件("faces_list.pkl")
if os.path.exists(ENCODINGS_FILE) and os.path.exists("faces_list.pkl"):
    # 如果文件存在，则从"faces_list.pkl"加载之前保存的人脸图片列表
    with open("faces_list.pkl", 'rb') as f:
        saved_faces_list = pickle.load(f)
        
    # 比较当前目录下的人脸图片列表(current_faces_list)与加载的保存列表(saved_faces_list)
    if current_faces_list == saved_faces_list:
        # 如果列表没有变化，直接从编码文件(ENCODINGS_FILE)加载已知人脸的编码和姓名
        known_faces, known_names = load_encodings(ENCODINGS_FILE)
    else:
        # 如果图片列表有变化（新增或删除了图片），则重新加载所有已知人脸的编码和姓名
        known_faces, known_names = load_known_faces(known_faces_dir)
        # 保存新的人脸编码和姓名到ENCODINGS_FILE
        save_encodings(known_faces, known_names, ENCODINGS_FILE)
        # 更新"faces_list.pkl"，保存当前的人脸图片列表
        with open("faces_list.pkl", 'wb') as f:
            pickle.dump(current_faces_list, f)
else:
    # 如果人脸编码文件或人脸图片列表文件不存在，则初始化过程
    # 加载或计算所有已知人脸的编码和姓名
    known_faces, known_names = load_known_faces(known_faces_dir)
    # 保存这些人脸数据到ENCODINGS_FILE
    save_encodings(known_faces, known_names, ENCODINGS_FILE)
    # 创建并保存当前的人脸图片列表到"faces_list.pkl"
    with open("faces_list.pkl", 'wb') as f:
        pickle.dump(current_faces_list, f)

# 启动图形化界面
create_gui()
